from langchain_core.tools import tool
from pydantic import BaseModel, Field
from typing import List, Dict, Optional
from agent.tools.database_connection import DatabaseConnection  # 导入数据库连接管理类

# 定义查询指标定义的输入模型
class IndicatorDefinitionQuery(BaseModel):
    indicator_id: int = Field(description="指标ID")

# 查询指标定义工具
@tool("query_indicator_definition", args_schema=IndicatorDefinitionQuery, return_direct=True)
async def query_indicator_definition(indicator_id: int) -> Dict:
    """根据指标ID查询指标定义，包括计算口径、下钻维度等"""
    try:
        # 使用数据库连接管理类获取数据库连接
        db_connection = DatabaseConnection.get_connection()
        if not db_connection:
            raise ValueError("Unable to connect to the database.")

        cursor = db_connection.cursor(dictionary=True)

        # 查询指标的基本信息
        cursor.execute(f"SELECT * FROM indicators WHERE id = {indicator_id}")
        indicator = cursor.fetchone()

        if not indicator:
            raise ValueError(f"Indicator with ID {indicator_id} not found.")

        # 查询下钻维度
        cursor.execute(f"SELECT dimension_id FROM indicator_breakdowns WHERE indicator_id = {indicator_id}")
        breakdowns = cursor.fetchall()

        dimensions = []
        for breakdown in breakdowns:
            cursor.execute(f"SELECT name FROM dimensions WHERE id = {breakdown['dimension_id']}")
            dimension = cursor.fetchone()
            if dimension:
                dimensions.append(dimension['name'])

        cursor.close()
        DatabaseConnection.close_connection(db_connection)

        # 构建指标定义的返回值
        indicator_definition = {
            "computation": indicator["computation"],
            "data_source": indicator["data_source"],
            "filters": indicator.get("filters", ""),
            "breakdowns": dimensions
        }

        return indicator_definition

    except Exception as e:
        print(f"Error retrieving indicator definition: {e}")
        return {}
